{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import astropy.cosmology as cosmo\n",
    "import astropy.units as u"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Produce comoving distances and distance moduli for cosmologies with multiple non-degenerate massive neutrinos, using astropy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the cosmological parameter we will use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Use these for each test\n",
    "OmegaC   = 0.25\n",
    "OmegaB   = 0.05\n",
    "h        = 0.7\n",
    "As       = 2.1*10**(-9)\n",
    "\n",
    "# We set the total number of neutrinos to 3 instead of 3.046\n",
    "# This is just because astropy only allows splitting Neff equally between all neutrinos species\n",
    "# So we need an integer to match our code.\n",
    "\n",
    "Nnu_tot  = 3  \n",
    "\n",
    "# Use 5 different sets of these parameters\n",
    "OmegaL    = [0.7, 0.7, 0.7, 0.65, 0.75]\n",
    "w0        = [-1., -0.9, -0.9, -0.9, -0.9]\n",
    "wa        = [0., 0., 0.1, 0.1, 0.1]\n",
    "\n",
    "mnu_set_0 = [0.04, 0.,   0.] * u.eV\n",
    "mnu_set_1 = [0.05, 0.01, 0.] * u.eV\n",
    "mnu_set_2 = [0.03, 0.02, 0.04] * u.eV\n",
    "mnu_set_3 = [0.05, 0.,   0.] * u.eV\n",
    "mnu_set_4 = [0.03, 0.02, 0.] * u.eV\n",
    "\n",
    "mnu = [mnu_set_0, mnu_set_1, mnu_set_2, mnu_set_3, mnu_set_4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up the 5 cosmologies with astropy cosmology objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cosmo0 = cosmo.w0waCDM(100*h, OmegaC + OmegaB, OmegaL[0], w0 = w0[0], wa = wa[0], Tcmb0 = 2.725, Neff = Nnu_tot, m_nu= mnu_set_0, Ob0 = OmegaB)\n",
    "cosmo1 = cosmo.w0waCDM(100*h, OmegaC + OmegaB, OmegaL[1], w0 = w0[1], wa = wa[1], Tcmb0 = 2.725, Neff = Nnu_tot, m_nu= mnu_set_1, Ob0 = OmegaB)\n",
    "cosmo2 = cosmo.w0waCDM(100*h, OmegaC + OmegaB, OmegaL[2], w0 = w0[2], wa = wa[2], Tcmb0 = 2.725, Neff = Nnu_tot, m_nu= mnu_set_2, Ob0 = OmegaB)\n",
    "cosmo3 = cosmo.w0waCDM(100*h, OmegaC + OmegaB, OmegaL[3], w0 = w0[3], wa = wa[3], Tcmb0 = 2.725, Neff = Nnu_tot, m_nu= mnu_set_3, Ob0 = OmegaB)\n",
    "cosmo4 = cosmo.w0waCDM(100*h, OmegaC + OmegaB, OmegaL[4], w0 = w0[4], wa = wa[4], Tcmb0 = 2.725, Neff = Nnu_tot, m_nu= mnu_set_4,Ob0 = OmegaB)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the redshifts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "z = [1., 2., 3., 4., 5.]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the comoving distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Loop over models\n",
    "chi0 = (cosmo0.comoving_distance(z)).value\n",
    "chi1 = (cosmo1.comoving_distance(z)).value\n",
    "chi2 = (cosmo2.comoving_distance(z)).value\n",
    "chi3 = (cosmo3.comoving_distance(z)).value\n",
    "chi4 = (cosmo4.comoving_distance(z)).value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Output these to a benchmark file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('./benchmark/chi_mnu_model1-5.txt', 'a') as out_file:\n",
    "    out_file.write('#z chi(z,model1) chi(z,model2) chi(z,model3) chi(z,model4) chi(z,model5)\\n')\n",
    "    out_file.write(str(z[0])+'\\t'+str(chi0[0])+'\\t'+str(chi1[0])+'\\t'+str(chi2[0])+'\\t'+str(chi3[0])+'\\t'+str(chi4[0])+'\\n')\n",
    "    out_file.write(str(z[1])+'\\t'+str(chi0[1])+'\\t'+str(chi1[1])+'\\t'+str(chi2[1])+'\\t'+str(chi3[1])+'\\t'+str(chi4[1])+'\\n')\n",
    "    out_file.write(str(z[2])+'\\t'+str(chi0[2])+'\\t'+str(chi1[2])+'\\t'+str(chi2[2])+'\\t'+str(chi3[2])+'\\t'+str(chi4[2])+'\\n')\n",
    "    out_file.write(str(z[3])+'\\t'+str(chi0[3])+'\\t'+str(chi1[3])+'\\t'+str(chi2[3])+'\\t'+str(chi3[3])+'\\t'+str(chi4[3])+'\\n')\n",
    "    out_file.write(str(z[4])+'\\t'+str(chi0[4])+'\\t'+str(chi1[4])+'\\t'+str(chi2[4])+'\\t'+str(chi3[4])+'\\t'+str(chi4[4])+'\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get also the distance moduli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dm0 = (cosmo0.distmod(z)).value\n",
    "dm1 = (cosmo1.distmod(z)).value\n",
    "dm2 = (cosmo2.distmod(z)).value\n",
    "dm3 = (cosmo3.distmod(z)).value\n",
    "dm4 = (cosmo4.distmod(z)).value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And output that to file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('./benchmark/dm_mnu_model1-5.txt', 'a') as out_file:\n",
    "    out_file.write('#z dm(z,model1) dm(z,model2) dm(z,model3) dm(z,model4) dm(z,model5)\\n')\n",
    "    out_file.write(str(z[0])+'\\t'+str(dm0[0])+'\\t'+str(dm1[0])+'\\t'+str(dm2[0])+'\\t'+str(dm3[0])+'\\t'+str(dm4[0])+'\\n')\n",
    "    out_file.write(str(z[1])+'\\t'+str(dm0[1])+'\\t'+str(dm1[1])+'\\t'+str(dm2[1])+'\\t'+str(dm3[1])+'\\t'+str(dm4[1])+'\\n')\n",
    "    out_file.write(str(z[2])+'\\t'+str(dm0[2])+'\\t'+str(dm1[2])+'\\t'+str(dm2[2])+'\\t'+str(dm3[2])+'\\t'+str(dm4[2])+'\\n')\n",
    "    out_file.write(str(z[3])+'\\t'+str(dm0[3])+'\\t'+str(dm1[3])+'\\t'+str(dm2[3])+'\\t'+str(dm3[3])+'\\t'+str(dm4[3])+'\\n')\n",
    "    out_file.write(str(z[4])+'\\t'+str(dm0[4])+'\\t'+str(dm1[4])+'\\t'+str(dm2[4])+'\\t'+str(dm3[4])+'\\t'+str(dm4[4])+'\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Now do this for high redshifts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "z_high = [10., 20., 50., 100., 200., 500., 1000.]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loop over models\n",
    "chi0_high = (cosmo0.comoving_distance(z_high)).value\n",
    "chi1_high = (cosmo1.comoving_distance(z_high)).value\n",
    "chi2_high = (cosmo2.comoving_distance(z_high)).value\n",
    "chi3_high = (cosmo3.comoving_distance(z_high)).value\n",
    "chi4_high = (cosmo4.comoving_distance(z_high)).value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('./benchmark/chi_hiz_mnu_model1-5.txt', 'a') as out_file:\n",
    "    out_file.write('#z chi(z,model1) chi(z,model2) chi(z,model3) chi(z,model4) chi(z,model5)\\n')\n",
    "    out_file.write(str(z_high[0])+'\\t'+str(chi0_high[0])+'\\t'+str(chi1_high[0])+'\\t'+str(chi2_high[0])+'\\t'+str(chi3_high[0])+'\\t'+str(chi4_high[0])+'\\n')\n",
    "    out_file.write(str(z_high[1])+'\\t'+str(chi0_high[1])+'\\t'+str(chi1_high[1])+'\\t'+str(chi2_high[1])+'\\t'+str(chi3_high[1])+'\\t'+str(chi4_high[1])+'\\n')\n",
    "    out_file.write(str(z_high[2])+'\\t'+str(chi0_high[2])+'\\t'+str(chi1_high[2])+'\\t'+str(chi2_high[2])+'\\t'+str(chi3_high[2])+'\\t'+str(chi4_high[2])+'\\n')\n",
    "    out_file.write(str(z_high[3])+'\\t'+str(chi0_high[3])+'\\t'+str(chi1_high[3])+'\\t'+str(chi2_high[3])+'\\t'+str(chi3_high[3])+'\\t'+str(chi4_high[3])+'\\n')\n",
    "    out_file.write(str(z_high[4])+'\\t'+str(chi0_high[4])+'\\t'+str(chi1_high[4])+'\\t'+str(chi2_high[4])+'\\t'+str(chi3_high[4])+'\\t'+str(chi4_high[4])+'\\n')\n",
    "    out_file.write(str(z_high[5])+'\\t'+str(chi0_high[5])+'\\t'+str(chi1_high[5])+'\\t'+str(chi2_high[5])+'\\t'+str(chi3_high[5])+'\\t'+str(chi4_high[5])+'\\n')\n",
    "    out_file.write(str(z_high[6])+'\\t'+str(chi0_high[6])+'\\t'+str(chi1_high[6])+'\\t'+str(chi2_high[6])+'\\t'+str(chi3_high[6])+'\\t'+str(chi4_high[6])+'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dm0_high = (cosmo0.distmod(z_high)).value\n",
    "dm1_high = (cosmo1.distmod(z_high)).value\n",
    "dm2_high = (cosmo2.distmod(z_high)).value\n",
    "dm3_high = (cosmo3.distmod(z_high)).value\n",
    "dm4_high = (cosmo4.distmod(z_high)).value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('./benchmark/dm_hiz_mnu_model1-5.txt', 'a') as out_file:\n",
    "    out_file.write('#z dm(z,model1) dm(z,model2) dm(z,model3) dm(z,model4) dm(z,model5)\\n')\n",
    "    out_file.write(str(z_high[0])+'\\t'+str(dm0_high[0])+'\\t'+str(dm1_high[0])+'\\t'+str(dm2_high[0])+'\\t'+str(dm3_high[0])+'\\t'+str(dm4_high[0])+'\\n')\n",
    "    out_file.write(str(z_high[1])+'\\t'+str(dm0_high[1])+'\\t'+str(dm1_high[1])+'\\t'+str(dm2_high[1])+'\\t'+str(dm3_high[1])+'\\t'+str(dm4_high[1])+'\\n')\n",
    "    out_file.write(str(z_high[2])+'\\t'+str(dm0_high[2])+'\\t'+str(dm1_high[2])+'\\t'+str(dm2_high[2])+'\\t'+str(dm3_high[2])+'\\t'+str(dm4_high[2])+'\\n')\n",
    "    out_file.write(str(z_high[3])+'\\t'+str(dm0_high[3])+'\\t'+str(dm1_high[3])+'\\t'+str(dm2_high[3])+'\\t'+str(dm3_high[3])+'\\t'+str(dm4_high[3])+'\\n')\n",
    "    out_file.write(str(z_high[4])+'\\t'+str(dm0_high[4])+'\\t'+str(dm1_high[4])+'\\t'+str(dm2_high[4])+'\\t'+str(dm3_high[4])+'\\t'+str(dm4_high[4])+'\\n')\n",
    "    out_file.write(str(z_high[5])+'\\t'+str(dm0_high[5])+'\\t'+str(dm1_high[5])+'\\t'+str(dm2_high[5])+'\\t'+str(dm3_high[5])+'\\t'+str(dm4_high[5])+'\\n')\n",
    "    out_file.write(str(z_high[6])+'\\t'+str(dm0_high[6])+'\\t'+str(dm1_high[6])+'\\t'+str(dm2_high[6])+'\\t'+str(dm3_high[6])+'\\t'+str(dm4_high[6])+'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
